import cv2
import numpy as np

SQRT_0_5 = 0.70710678118654757

class AutoWire():
    """
    Asimple livewire implementation for verification using
        1.Canny edge detector + gradient magnitude + gradient direction
        2.Dijkstra algorithm
    """
    def __init__(self, image):

        self.img = image
        self.x_lim, self.y_lim =   self.img.shape[0:2]

        # canny 算法获取边缘 1减去大小反转
        self.cost_edges = 1 - cv2.Canny( self.img, 85, 170) / 255.0
        self.grad_x, self.grad_y, self.grad_mag = self._get_grad( self.img)
        self.cost_grad_mag = 1 - self.grad_mag / np.max(self.grad_mag)

        self.weight = (0.425, 0.425, 0.15)
        self.n_pixs = self.x_lim * self.y_lim
        self.n_processed = 0

        self.x0=0
        self.y0=0
        self.x1=0
        self.y1 = 0


    def setregion(self,x0,y0,x1,y1):
        self.x0 = x0
        self.y0 = y0
        self.x1 = x1
        self.y1 = y1

    @classmethod
    def _get_grad(cls, image):
        """
        Returnthe gradient magnitude of the image using Sobel operator
        """
        rgb = True if len(image.shape) > 2 else False
        grad_x = cv2.Sobel(image, cv2.CV_64F, 1, 0, ksize=3)
        grad_y = cv2.Sobel(image, cv2.CV_64F, 0, 1, ksize=3)
        if rgb:
            # A very rough approximation for quick verification...
            grad_x = np.max(grad_x, axis=2)
            grad_y = np.max(grad_y, axis=2)

        grad_mag = np.sqrt(grad_x ** 2 + grad_y ** 2)
        grad_x /= grad_mag
        grad_y /= grad_mag

        return grad_x, grad_y, grad_mag

    def _get_neighbors(self, p):
        """
        Return8 neighbors around the pixel p
        """
        #定义局部的搜索点
        tx1=min(self.x0,self.x1)
        ty1=min(self.y0,self.y1)

        tx2 = max(self.x0, self.x1)
        ty2 = max(self.y0, self.y1)

        x, y = p
        x0 = tx1 if x <= tx1 else x - 1
        x1 = tx2 if x >= tx2 - 1 else x + 2
        y0 = ty1 if y <= ty1 else y - 1
        y1 = ty2 if y >=ty2 - 1 else y + 2

        if(x1==tx2):
            x1+=1
        if(x0==x1):
            x1+=1
        if (y1 == ty2):
            y1 += 1
        if(y0==y1):
            y1+=1
        return [(x, y) for x in range(x0, x1) for y in range(y0, y1) if (x, y) != p]

    def _get_grad_direction_cost(self, p, q):
        """
        Calculatethe gradient changes refer to the link direction
        """
        dp = (self.grad_y[p[0]][p[1]], -self.grad_x[p[0]][p[1]])
        dq = (self.grad_y[q[0]][q[1]], -self.grad_x[q[0]][q[1]])

        l = np.array([q[0] - p[0], q[1] - p[1]], np.float)
        if 0 not in l:
            l *= SQRT_0_5

        dp_l = np.dot(dp, l)
        l_dq = np.dot(l, dq)
        if dp_l < 0:
            dp_l = -dp_l
            l_dq = -l_dq

        # 2/3pi * ...
        return 0 # 0.212206590789 * (np.arccos(dp_l) + np.arccos(l_dq))

    def _local_cost(self, p, q):
        """
        1.Calculate the Canny edges & gradient magnitude cost taking into accountEuclidean distance
        2.Combine with gradient direction
        Assumption:p & q are neighbors
        """
        diagnol = q[0] == p[0] or q[1] == p[1]

        # c0, c1 and c2 are costs from Canny operator, gradient magnitude and gradientdirection respectively
        if diagnol:
            c0 = self.cost_edges[p[0]][p[1]] - SQRT_0_5 * (self.cost_edges[p[0]][p[1]] - self.cost_edges[q[0]][q[1]])
            c1 = self.cost_grad_mag[p[0]][p[1]] - SQRT_0_5 * (self.cost_grad_mag[p[0]][p[1]] - self.cost_grad_mag[q[0]][q[1]])
            c2 = SQRT_0_5 * self._get_grad_direction_cost(p, q)
        else:
            c0 = self.cost_edges[q[0]][q[1]]
            c1 = self.cost_grad_mag[q[0]][q[1]]
            c2 = self._get_grad_direction_cost(p, q)

        if np.isnan(c2):
            c2 = 0.0

        w0, w1, w2 = self.weight
        cost_pq = w0 * c0 + w1 * c1 + w2 * c2

        return cost_pq * cost_pq

    def get_path_matrix(self, seed):
        """
        Getthe back tracking matrix of the whole image from the cost matrix
        """
        try:
            neighbors = []  # 8 neighbors of thepixel being processed
            processed = set()  # Processed point
            cost = {seed: 0.0}  # Accumulated cost, initializedwith seed to itself
            paths = {}

            self.n_processed = 0

            while cost:
                # Expand the minimum cost point
                p = min(cost, key=cost.get)
                neighbors = self._get_neighbors(p)
                processed.add(p)

                # Record accumulated costs and back tracking point for newly expanded points
                for q in [x for x in neighbors if x not in processed]:
                    temp_cost = cost[p] + self._local_cost(p, q)
                    if q in cost:
                        if temp_cost < cost[q]:
                            cost.pop(q)
                            paths[q] = p
                    else:
                        cost[q] = temp_cost
                        processed.add(q)
                        paths[q] = p

                    # Pop traversed points
                cost.pop(p)
                self.n_processed += 1
                if(self.n_processed>=self.n_pixs):
                    break
        except Exception as e:
            print(e)

        return paths
